{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "### nq2squad.py\n",
    "把源竞赛数据集`simplified-nq-t*.jsonl`文件转换成NQ json的格式"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 导入依赖包"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import argparse\n",
    "import pandas as pd\n",
    "from tqdm import tqdm"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### SQuAD 2.0格式分析"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Collected by Gao Xiangfu\n",
    "\n",
    "\n",
    "- 来自 `dev-v2.0/json`\n",
    "\n",
    "### 顶层数据\n",
    "\n",
    "```JavaScript\n",
    "{\n",
    "    \"version\": \"2.0\",\n",
    "    \"data\": [\n",
    "        // 包含 34 个数据字典，独有标题，包含多个文段及附属问题\n",
    "    ]\n",
    "}\n",
    "```\n",
    "\n",
    "### 数据字典\n",
    "\n",
    "```JavaScript\n",
    "{\n",
    "    \"title\": \"标题\",\n",
    "    \"paragraphs\": [\n",
    "        // 多个阅读理解\n",
    "        ...,\n",
    "        {\n",
    "            // 每个文段问题数不定量\n",
    "            \"qas\": [\n",
    "                ...,\n",
    "                {\n",
    "                    \"question\": \"(str)\",\n",
    "                    \"id\": \"(str)\",\n",
    "                    \"answers\": [\n",
    "                        ...,\n",
    "                        {\n",
    "                            \"text\": \"(str)\",\n",
    "                            \"answer_start\": \"按char计的开始位置(int)\"\n",
    "                        },\n",
    "                        ...\n",
    "                    ],\n",
    "                    \"is_impossible\": \"是否可以回答(bool)\",\n",
    "                },\n",
    "                ...\n",
    "            ]\n",
    "            \"context\": \"文段内容(str)\"\n",
    "        },\n",
    "        ...\n",
    "    ]\n",
    "\n",
    "}\n",
    "```\n",
    "\n",
    "### Original Natural Questions\n",
    "\n",
    "\n",
    "每一行是一个样例字典。针对一个问题`question_text`字符串格式，有`question_tokens`可用。\n",
    "\n",
    "`document_tokens` 为主要使用的阅读理解文章来源。针对 `document_html` 的标记经实验验证有大量偏移错误。可能难以使用。 \n",
    "\n",
    "### 标记 annotations\n",
    "\n",
    "每个问题可能有多个答案标记。\n",
    "\n",
    "```JavaScript\n",
    "{\n",
    "    \"annotations\": [\n",
    "        ...,\n",
    "        {\n",
    "            \"annotations_id\": \"(int)\",\n",
    "            \"long_answer\": {\n",
    "                \"candidate_index\": \"对应样例顶级long_answer_candidates，不存在为 -1\",\n",
    "                \"start_token\": \"对应顶级document_tokens，左闭右开\"\n",
    "                \"end_token\": \"同上\"\n",
    "                \"start_byte\": \"用不上\",\n",
    "                \"end_byte\": \"用不上\"\n",
    "            }\n",
    "            \"short_answers\": [\"可能有多个短答案，格式同长答案\"],\n",
    "            \"yes_no_answer\": \"NONE or YES or NO\"\n",
    "        }\n",
    "        ...,\n",
    "    ]\n",
    "}\n",
    "\n",
    "```\n",
    "\n",
    "其他数据分析参考[Github - Natural Questions](https://github.com/google-research-datasets/natural-questions)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### convert_nq_to_squad函数(add comments）"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def convert_nq_to_squad(args=None):\n",
    "    np.random.seed(123)\n",
    "    if args is None:\n",
    "        # 添加配置参数\n",
    "        parser = argparse.ArgumentParser()\n",
    "        parser.add_argument('--fn', type=str,\n",
    "                            default=r'F:\\dataset_download\\nqa2_kaggle_competition\\simplified-nq-train.jsonl')\n",
    "        parser.add_argument('--version', type=str, default='v1.0.1')\n",
    "        parser.add_argument('--prefix', type=str, default='nq')\n",
    "        parser.add_argument('--p_val', type=float, default=0.1)\n",
    "        parser.add_argument('--crop_len', type=int, default=2_500)\n",
    "        parser.add_argument('--num_samples', type=int, default=100_000)\n",
    "        parser.add_argument('--val_ids', type=str, default='val_ids.csv')\n",
    "        parser.add_argument('--do_enumerate', action='store_true')\n",
    "        parser.add_argument('--do_not_dump', action='store_true')\n",
    "        parser.add_argument('--num_max_tokens', type=int, default=400_000)\n",
    "        args = parser.parse_args()\n",
    "        \n",
    "    # 根据加载的数据类型决定转换出的文件名\n",
    "    is_train = 'train' in args.fn\n",
    "    if is_train:\n",
    "        train_fn = f'{args.prefix}-train-{args.version}.json'\n",
    "        val_fn = f'{args.prefix}-val-{args.version}.json'\n",
    "        print(f'Converting {args.fn} to {train_fn} & {val_fn} ... ')\n",
    "    else:\n",
    "        test_fn = f'{args.prefix}-test-{args.version}.json'\n",
    "        print(f'Converting {args.fn} to {test_fn} ... ')\n",
    "\n",
    "    if args.val_ids:\n",
    "        val_ids = set(str(x) for x in pd.read_csv(args.val_ids)['val_ids'].values)\n",
    "    else:\n",
    "        val_ids = set()\n",
    "        \n",
    "    # 定义一些变量\n",
    "    entries = []\n",
    "    smooth = 0.999\n",
    "    total_split_len, long_split_len = 0., 0.\n",
    "    long_end = 0.\n",
    "    num_very_long, num_yes_no, num_short_dropped, num_trimmed = 0, 0, 0, 0\n",
    "    num_short_possible, num_long_possible = 0, 0\n",
    "    max_end_token = -1\n",
    "    orig_data = {} # 原始数据——dict形式\n",
    "    \n",
    "    with open(args.fn) as f: # 打开要转换的文件\n",
    "        progress = tqdm(f, total=args.num_samples) # 加载进度条\n",
    "        entry = {}\n",
    "        for kk, line in enumerate(progress): # 获取每行json数据和行下标\n",
    "            if kk >= args.num_samples:  # 若当前处理的行数已达到参数规定的数目，则跳出\n",
    "                break\n",
    "\n",
    "            data = json.loads(line) # 解码json对象, 转换为dict形式\n",
    "            data_cpy = data.copy() # 浅复制，不完全独立\n",
    "            example_id = str(data_cpy.pop('example_id')) # 从字典数据中pop出键example_id对应的值并暂存\n",
    "            data_cpy['document_text'] = '' # 键document_text对应的值置空，不影响data\n",
    "            orig_data[example_id] = data_cpy # 往字典中添加键值对<example_id, data_cpy>\n",
    "            url = 'MISSING' if not is_train else data['document_url']\n",
    "            # progress.write(f'############ {url} ###############')\n",
    "            document_text = data['document_text'] # 获取文本\n",
    "            document_text_split = document_text.split(' ') # 文本按空格进行划分\n",
    "            # trim super long\n",
    "            if len(document_text_split) > args.num_max_tokens:\n",
    "                num_trimmed += 1\n",
    "                document_text_split = document_text_split[:args.num_max_tokens] # 修剪文本\n",
    "\n",
    "            if args.do_enumerate:\n",
    "                document_text_split = enumerate_tags(document_text_split) # 进行\n",
    "            question = data['question_text']  # + '?'\n",
    "            annotations = [None] if not is_train else data['annotations'] # 获取答案annotation\n",
    "            assert len(annotations) == 1, annotations\n",
    "            # User str keys!\n",
    "            example_id = str(data['example_id'])\n",
    "            candidates = data['long_answer_candidates'] # 侯选的长答案列表\n",
    "            if not is_train:\n",
    "                qa = {'question': question, 'id': example_id, 'crop_start': 0}\n",
    "                context = ' '.join(document_text_split)\n",
    "            \n",
    "            # 若加载的是训练集\n",
    "            else:\n",
    "                long_answer = annotations[0]['long_answer'] # 取出长答案\n",
    "                long_answer_len = long_answer['end_token'] - long_answer['start_token']\n",
    "                total_split_len = smooth * total_split_len + (1. - smooth) * len(\n",
    "                    document_text_split)\n",
    "                long_split_len = smooth * long_split_len + (1. - smooth) * \\\n",
    "                                 long_answer_len\n",
    "                if long_answer['end_token'] > 0:\n",
    "                    long_end = smooth * long_end + (1. - smooth) * long_answer['end_token']\n",
    "\n",
    "                if long_answer['end_token'] > max_end_token:\n",
    "                    max_end_token = long_answer['end_token']\n",
    "\n",
    "                progress.set_postfix({'ltotal': int(total_split_len),\n",
    "                                      'llong': int(long_split_len), 'long_end': round(long_end, 2)})\n",
    "\n",
    "                short_answers = annotations[0]['short_answers'] # 取出短答案\n",
    "                yes_no_answer = annotations[0]['yes_no_answer'] # 取出yes_no_answer\n",
    "                # 跳过回答yes/no/none的问题\n",
    "                if yes_no_answer != 'NONE':\n",
    "                    # progress.write(f'Skipping yes-no: {yes_no_answer}')\n",
    "                    num_yes_no += 1\n",
    "                    continue\n",
    "\n",
    "                # print(f'Q: {question}')\n",
    "                # print(f'L: {long_answer_str}')\n",
    "                long_is_impossible = long_answer['start_token'] == -1 # 将long_is_impossible的bool值同long_answer的起始token取值情况联系起来\n",
    "                if long_is_impossible:\n",
    "                    long_answer_candidate = np.random.randint(len(candidates))  # 随机选一个候选的长答案\n",
    "                else:\n",
    "                    long_answer_candidate = long_answer['candidate_index'] # 否则选定正确答案\n",
    "\n",
    "                long_start_token = candidates[long_answer_candidate]['start_token']\n",
    "                long_end_token = candidates[long_answer_candidate]['end_token']\n",
    "                # generate crop based on tokens. Note that validation samples should\n",
    "                # not be cropped as this won't reflect test set performance.\n",
    "                if args.crop_len > 0 and example_id not in val_ids:\n",
    "                    crop_start = long_start_token - np.random.randint(int(args.crop_len * 0.75))\n",
    "                    if crop_start <= 0:\n",
    "                        crop_start = 0\n",
    "                        crop_start_len = -1\n",
    "                    else:\n",
    "                        crop_start_len = len(' '.join(document_text_split[:crop_start]))\n",
    "\n",
    "                    crop_end = crop_start + args.crop_len\n",
    "                else:\n",
    "                    crop_start = 0\n",
    "                    crop_start_len = -1\n",
    "                    crop_end = 10_000_000\n",
    "\n",
    "                is_very_long = False\n",
    "                if long_end_token > crop_end:\n",
    "                    num_very_long += 1\n",
    "                    is_very_long = True\n",
    "                    # progress.write(f'{num_very_long}: Skipping very long answer {long_end_token}, {crop_end}')\n",
    "                    # continue\n",
    "                # 裁剪document_text\n",
    "                document_text_crop_split = document_text_split[crop_start: crop_end]\n",
    "                context = ' '.join(document_text_crop_split)\n",
    "                # 构建长答案\n",
    "                long_answers_ = []\n",
    "                if not long_is_impossible:\n",
    "                    long_answer_pre_split = document_text_split[:long_answer[\n",
    "                        'start_token']]\n",
    "                    long_answer_start = len(' '.join(long_answer_pre_split)) - \\\n",
    "                                        crop_start_len\n",
    "                    long_answer_split = document_text_split[long_answer['start_token']:\n",
    "                                                            long_answer['end_token']]\n",
    "                    \n",
    "                    # 把从document_text中把长答案对应的文本裁出来并转成字符串\n",
    "                    long_answer_text = ' '.join(long_answer_split)\n",
    "                    if not is_very_long:\n",
    "                        assert context[long_answer_start: long_answer_start + len(\n",
    "                            long_answer_text)] == long_answer_text, long_answer_text\n",
    "                        \n",
    "                    # 长答案构建完毕\n",
    "                    long_answers_ = [{'text': long_answer_text,\n",
    "                                      'answer_start': long_answer_start}]\n",
    "\n",
    "                # 构建短答案\n",
    "                short_is_impossible = len(short_answers) == 0\n",
    "                short_answers_ = []\n",
    "                if not short_is_impossible:\n",
    "                    for short_answer in short_answers:\n",
    "                        short_start_token = short_answer['start_token']\n",
    "                        short_end_token = short_answer['end_token']\n",
    "                        if short_start_token >= crop_start + args.crop_len:\n",
    "                            num_short_dropped += 1\n",
    "                            continue\n",
    "                        short_answers_pre_split = document_text_split[:short_start_token]\n",
    "                        short_answer_start = len(' '.join(short_answers_pre_split)) - \\\n",
    "                                             crop_start_len\n",
    "                        short_answer_split = document_text_split[short_start_token: short_end_token]\n",
    "                        # 从document_text中裁出短答案并转成字符串形式\n",
    "                        short_answer_text = ' '.join(short_answer_split)\n",
    "                        assert short_answer_text != ''\n",
    "\n",
    "                        # this happens if we crop and parts of the short answer overflow\n",
    "                        short_from_context = context[short_answer_start: short_answer_start + len(short_answer_text)]\n",
    "                        if short_from_context != short_answer_text:\n",
    "                            print(f'short diff: {short_from_context} vs {short_answer_text}')\n",
    "                        # 短答案构建完成\n",
    "                        short_answers_.append({'text': short_from_context,\n",
    "                                               'answer_start': short_answer_start})\n",
    "\n",
    "                if len(short_answers_) == 0:\n",
    "                    short_is_impossible = True\n",
    "\n",
    "                if not short_is_impossible:\n",
    "                    num_short_possible += 1\n",
    "                if not long_is_impossible:\n",
    "                    num_long_possible += 1\n",
    "                # 转成squad的形式\n",
    "                qa = {'question': question,\n",
    "                      'short_answers': short_answers_, 'long_answers': long_answers_,\n",
    "                      'id': example_id, 'short_is_impossible': short_is_impossible,\n",
    "                      'long_is_impossible': long_is_impossible,\n",
    "                      'crop_start': crop_start}\n",
    "\n",
    "            paragraph = {'qas': [qa], 'context': context}\n",
    "            entry = {'title': url, 'paragraphs': [paragraph]}\n",
    "            entries.append(entry)\n",
    "\n",
    "    progress.write('  ------------ STATS ------------------')\n",
    "    progress.write(f'  Found {num_yes_no} yes/no, {num_very_long} very long'\n",
    "                   f' and {num_short_dropped} short of {kk} and trimmed {num_trimmed}')\n",
    "    progress.write(f'  #short {num_short_possible} #long {num_long_possible}'\n",
    "                   f' of {len(entries)}')\n",
    "\n",
    "    if is_train:\n",
    "        train_entries, val_entries = [], []\n",
    "        for entry in entries:\n",
    "            if entry['paragraphs'][0]['qas'][0]['id'] not in val_ids:\n",
    "                train_entries.append(entry)\n",
    "            else:\n",
    "                val_entries.append(entry)\n",
    "\n",
    "        for out_fn, entries in [(train_fn, train_entries), (val_fn, val_entries)]:\n",
    "            if not args.do_not_dump:\n",
    "                with open(out_fn, 'w') as f:\n",
    "                    json.dump({'version': args.version, 'data': entries}, f)\n",
    "                progress.write(f'Wrote {len(entries)} entries to {out_fn}')\n",
    "\n",
    "            # save val in competition csv format\n",
    "            if 'val' in out_fn:\n",
    "                val_example_ids, val_strs = [], []\n",
    "                for entry in entries:\n",
    "                    example_id = entry['paragraphs'][0]['qas'][0]['id']\n",
    "                    short_answers = orig_data[example_id]['annotations'][0][\n",
    "                        'short_answers']\n",
    "                    sa_str = ''\n",
    "                    for si, sa in enumerate(short_answers):\n",
    "                        sa_str += f'{sa[\"start_token\"]}:{sa[\"end_token\"]}'\n",
    "                        if si < len(short_answers) - 1:\n",
    "                            sa_str += ' '\n",
    "                    val_example_ids.append(example_id + '_short')\n",
    "                    val_strs.append(sa_str)\n",
    "\n",
    "                    la = orig_data[example_id]['annotations'][0][\n",
    "                        'long_answer']\n",
    "                    la_str = ''\n",
    "                    if la['start_token'] > 0:\n",
    "                        la_str += f'{la[\"start_token\"]}:{la[\"end_token\"]}'\n",
    "                    val_example_ids.append(example_id + '_long')\n",
    "                    val_strs.append(la_str)\n",
    "\n",
    "                val_df = pd.DataFrame({'example_id': val_example_ids,\n",
    "                                       'PredictionString': val_strs})\n",
    "                val_csv_fn = f'{args.prefix}-val-{args.version}.csv'\n",
    "                val_df.to_csv(val_csv_fn, index=False, columns=['example_id',\n",
    "                                                                'PredictionString'])\n",
    "                print(f'Wrote csv to {val_csv_fn}')\n",
    "\n",
    "    else:\n",
    "        if not args.do_not_dump:\n",
    "            with open(test_fn, 'w') as f:\n",
    "                json.dump({'version': args.version, 'data': entries}, f)\n",
    "            progress.write(f'Wrote to {test_fn}')\n",
    "\n",
    "    if args.val_ids:\n",
    "        print(f'Using val ids from: {args.val_ids}')\n",
    "    return entries"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% \n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### SQuAD2.0数据模板中的各个数据如何产生？"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 顶层数据version及data\n",
    "- version\n",
    "\n",
    "      parser.add_argument('--version', type=str, default='v1.0.1')\n",
    "\n",
    "- data(包含许多数据字典)\n",
    "\n",
    "      entry = {'title': url, 'paragraphs': [paragraph]}\n",
    "      entries.append(entry)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 数据字典\n",
    "#### title\n",
    "\n",
    "      # url用作title\n",
    "      url = 'MISSING' if not is_train else data['document_url']\n",
    "    \n",
    "#### annotations\n",
    "- annotation\n",
    "\n",
    "      # 导入的每行数据都是一个json对象\n",
    "      data = json.loads(line)\n",
    "      annotations = [None] if not is_train else data['annotations']\n",
    "      \n",
    "- long_answer\n",
    "\n",
    "      long_answer = annotations[0]['long_answer']\n",
    "\n",
    "- short_answer\n",
    "\n",
    "      short_answers = annotations[0]['short_answers']\n",
    "      \n",
    "- yes_no_answer\n",
    "\n",
    "      yes_no_answer = annotations[0]['yes_no_answer']\n",
    "      # 跳过yes/no/none这一类的问题，并返回计数\n",
    "      if yes_no_answer != 'NONE':\n",
    "          # progress.write(f'Skipping yes-no: {yes_no_answer}')\n",
    "          num_yes_no += 1\n",
    "          continue\n",
    "          \n",
    "#### qa\n",
    "- question\n",
    "      \n",
    "      # 提取问题\n",
    "      question = data['question_text']\n",
    "\n",
    "- id\n",
    "\n",
    "      example_id = str(data['example_id'])\n",
    "\n",
    "- is_impossible\n",
    "\n",
    "      # 判断是否可回答（通过将long_is_impossible的bool值同long_answer的起始token取值情况联系起来）\n",
    "      long_is_impossible = long_answer['start_token'] == -1\n",
    "      # 通过将short_is_impossible的bool值同short_answer长度取值联系起来\n",
    "      short_is_impossible = len(short_answers) == 0\n",
    "    \n",
    "- answers(text及answer_start)\n",
    "    \n",
    "      # long_answers\n",
    "      # 如果可回答\n",
    "      if not long_is_impossible:\n",
    "          long_answer_pre_split = document_text_split[:long_answer[\n",
    "              'start_token']]\n",
    "          long_answer_start = len(' '.join(long_answer_pre_split)) - \\\n",
    "                              crop_start_len\n",
    "          long_answer_split = document_text_split[long_answer['start_token']:\n",
    "                                                  long_answer['end_token']]\n",
    "          long_answer_text = ' '.join(long_answer_split)\n",
    "          long_answers_ = [{'text': long_answer_text,\n",
    "                            'answer_start': long_answer_start}]\n",
    "                            \n",
    "      # short answers\n",
    "      short_answers_ = []\n",
    "      # 如果可回答\n",
    "      if not short_is_impossible:\n",
    "          for short_answer in short_answers:\n",
    "              short_start_token = short_answer['start_token']\n",
    "              short_end_token = short_answer['end_token']\n",
    "              if short_start_token >= crop_start + args.crop_len:\n",
    "                  num_short_dropped += 1\n",
    "                  continue\n",
    "              short_answers_pre_split = document_text_split[:short_start_token]\n",
    "              short_answer_start = len(' '.join(short_answers_pre_split)) - \\\n",
    "                                   crop_start_len\n",
    "              short_answer_split = document_text_split[short_start_token: short_end_token]\n",
    "              short_answer_text = ' '.join(short_answer_split)\n",
    "              assert short_answer_text != ''\n",
    "\n",
    "              # this happens if we crop and parts of the short answer overflow\n",
    "              short_from_context = context[short_answer_start: short_answer_start + len(short_answer_text)]\n",
    "              if short_from_context != short_answer_text:\n",
    "                  print(f'short diff: {short_from_context} vs {short_answer_text}')\n",
    "              short_answers_.append({'text': short_from_context,\n",
    "                                     'answer_start': short_answer_start})\n",
    "\n",
    "    \n",
    "#### paragraphs(qas及context)\n",
    "- qas([qa])\n",
    "      \n",
    "      # 将问题、回答、是否可答、id等放入字典\n",
    "      qa = {'question': question,\n",
    "          'short_answers': short_answers_, 'long_answers': long_answers_,\n",
    "          'id': example_id, 'short_is_impossible': short_is_impossible,\n",
    "          'long_is_impossible': long_is_impossible,\n",
    "          'crop_start': crop_start}\n",
    "          \n",
    "      paragraph = {'qas': [qa], 'context': context}\n",
    "    \n",
    "- context\n",
    "\n",
    "        # 取出源数据集中的document_text\n",
    "        document_text = data['document_text']\n",
    "        # 按空格划分文本\n",
    "        document_text_split = document_text.split(' ')\n",
    "        # 默认执行\n",
    "        if args.do_enumerate:\n",
    "            document_text_split = enumerate_tags(document_text_split)\n",
    "        # crop过后的文本\n",
    "        document_text_crop_split = document_text_split[crop_start: crop_end]\n",
    "        # 转成字符串\n",
    "        context = ' '.join(document_text_crop_split)\n",
    "    "
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### enumerate_tag()函数"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def enumerate_tags(text_split):\n",
    "  \"\"\"Reproduce the pre processing from:\n",
    "  A BERT Baseline for the Natural Questions (https://arxiv.org/pdf/1901.08634.pdf)\n",
    "\n",
    "  We introduce special markup tokens in the document  to  give  the  model\n",
    "  a  notion  of  which  part of the document it is reading.  The special\n",
    "  tokens we introduced are of the form “[Paragraph=N]”,“[Table=N]”, and “[List=N]”\n",
    "  at the beginning of the N-th paragraph,  list and table respectively\n",
    "  in the document. This decision was based on the observation that the first\n",
    "  few paragraphs and tables in the document are much more likely than the rest\n",
    "  of the document to contain the annotated answer and so the model could benefit\n",
    "  from knowing whether it is processing one of these passages.\n",
    "\n",
    "  We deviate as follows: Tokens are only created for the first 10 times. All other\n",
    "  tokens are the same. We only add `special_tokens`. These two are added as they\n",
    "  make 72.9% + 19.0% = 91.9% of long answers.\n",
    "  (https://github.com/google-research-datasets/natural-questions)\n",
    "  \"\"\"\n",
    "  special_tokens = ['<P>', '<Table>']\n",
    "  # 计数列表初始化\n",
    "  special_token_counts = [0 for _ in range(len(special_tokens))]\n",
    "  # 遍历按空格分割的document文本\n",
    "  for index, token in enumerate(text_split):\n",
    "    # 按不同special_token出现的次数进行计数\n",
    "    for special_token_index, special_token in enumerate(special_tokens):\n",
    "      if token == special_token:\n",
    "        cnt = special_token_counts[special_token_index]\n",
    "        # 只为不同special_token各自的前十项进行tag的修改\n",
    "        if cnt <= 10:\n",
    "          # 加上cnt就能指示读到文档的当前大概的位置\n",
    "          text_split[index] = f'<{special_token[1: -1]}{cnt}>'\n",
    "        special_token_counts[special_token_index] = cnt + 1\n",
    "\n",
    "  return text_split"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 引入特殊令牌\n",
    "\n",
    "- 特殊令牌\n",
    "\n",
    "      特殊令牌的作用：用于指示文档当前大概的位置\n",
    "      上述函数封装性很好，当special_token列表需要改变而需要达到\n",
    "      相同功能时，仅需要修改第一行即可。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "#### 读代码时不是很明白的一些点\n",
    "- 为什么要对document_text进行crop以及crop的规则?\n",
    "\n",
    "      # generate crop based on tokens. Note that validation samples should\n",
    "      # not be cropped as this won't reflect test set performance.\n",
    "      if args.crop_len > 0 and example_id not in val_ids:\n",
    "          crop_start = long_start_token - np.random.randint(int(args.crop_len * 0.75))\n",
    "          if crop_start <= 0:\n",
    "              crop_start = 0\n",
    "              crop_start_len = -1\n",
    "          else:\n",
    "              crop_start_len = len(' '.join(document_text_split[:crop_start]))\n",
    "\n",
    "              crop_end = crop_start + args.crop_len\n",
    "          else:\n",
    "              crop_start = 0\n",
    "              crop_start_len = -1\n",
    "              crop_end = 10_000_000\n",
    "              \n",
    "  \n",
    "- is_impossible的判断标准?\n",
    "\n",
    "      long_is_impossible = long_answer['start_token'] == -1\n",
    "      short_is_impossible = len(short_answers) == 0\n",
    "      \n",
    "- smooth的作用?\n",
    "\n",
    "      total_split_len = smooth * total_split_len + (1. - smooth) * len(\n",
    "                    document_text_split)\n",
    "      long_split_len = smooth * long_split_len + (1. - smooth) * \\\n",
    "                                 long_answer_len\n",
    "      if long_answer['end_token'] > 0:\n",
    "          long_end = smooth * long_end + (1. - smooth) * long_answer['end_token']\n",
    "\n",
    "      if long_answer['end_token'] > max_end_token:\n",
    "          max_end_token = long_answer['end_token']\n",
    "\n",
    "      progress.set_postfix({'ltotal': int(total_split_len),\n",
    "                            'llong': int(long_split_len), 'long_end': round(long_end, 2)})\n",
    "\n",
    "- val_ids.csv"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}